// @Author EthanScriptOn
// @Desc
package pool

import (
	"errors"
	"fmt"
	"gitee.com/fatzeng/srf_switch_basic_components/safe"
	"gitee.com/fatzeng/srf_switch_basic_components/tool"
	"sync"
	"sync/atomic"
	"time"
)

// WorkerPool The work pool, which accepts the task and selects a suitable worker to work on,
// is based on the thread pool in Java
type WorkerPool struct {
	coreSize            int                 // coreSize The number of cores in the work pool
	maxSize             int                 // maxSize The maximum number of working pools
	workerFactory       WorkerFactory       // workerFactory Creation of factories for workers
	nonCoreSurvivalTime time.Duration       // nonCoreSurvivalTime (Maximum number - number of cores) of the worker's survival time
	workerPool          map[Worker]struct{} // workerPool Work pools
	taskQueue           chan Task           // taskQueue Task queues
	rejectStrategy      RejectStrategy      // rejectStrategy Deny policy
	survivalNum         int32               // survivalNum Marks the number of surviving workers at the moment
	mu                  *sync.RWMutex       // mu Lock
	tp                  *sync.RWMutex       // tp Lock
	isClose             bool                // isClose Whether the working pool has been shut down
	workerSelection     WorkerSelection     // workerSelection Select a worker acquisition method
	safeGo              *safe.SafeGo        // safeGo With recovery mechanism goroutine
}

func GenerateWorkerPool(coreSize int, maxSize int, workerFactory WorkerFactory, nonCoreSurvivalTime time.Duration, rejectStrategy RejectStrategy, workerSelection WorkerSelection) *WorkerPool {
	return &WorkerPool{
		coreSize:            coreSize,
		maxSize:             maxSize,
		workerFactory:       workerFactory,
		nonCoreSurvivalTime: nonCoreSurvivalTime,
		workerPool:          make(map[Worker]struct{}),
		taskQueue:           make(chan Task),
		rejectStrategy:      rejectStrategy,
		mu:                  &sync.RWMutex{},
		tp:                  &sync.RWMutex{},
		workerSelection:     workerSelection,
		safeGo:              safe.GenerateSafeGo(),
	}
}

// Submit Submit a task
func (w *WorkerPool) Submit(task Task) (err error) {
	// The task cannot be empty
	if task == nil {
		return errors.New("task is nil")
	}
	return tool.LockWithSync(w.tp.RLock, w.tp.RUnlock, func() error {
		// Has the pool been closed? Turn off the post-and-go reject policy
		if w.IsClose() {
			w.rejectStrategy.Reject(task, w)
			return errors.New("worker pool is close")
		}
		// Determine whether to add a worker based on the number of survivors
		survivalNum := w.getSurvivalNum()
		if survivalNum < w.coreSize {
			if err = w.addWorker(task, true); err == nil {
				return nil
			}
		}
		// There is no need to add a worker to deliver the task to the work queue for slow processing,
		// and the work queue will not be out of order
		w.taskQueue <- task
		if err = w.addWorker(nil, false); err != nil {
			w.rejectStrategy.Reject(task, w)
		}
		return nil
	})
}

// IsClose Whether the working pool has been shut down
func (w *WorkerPool) IsClose() bool {
	var isClose bool
	_ = tool.LockWithSync(w.tp.RLock, w.tp.RUnlock, func() error {
		isClose = w.isClose
		return nil
	})
	return isClose
}

// Close Close a work pool
func (w *WorkerPool) Close() {
	_ = tool.LockWithSync(w.tp.Lock, w.tp.Unlock, func() error {
		w.isClose = true
		return nil
	})
}

// getSurvivalNum Obtain the surviving worker at the moment
func (w *WorkerPool) getSurvivalNum() int {
	return int(atomic.LoadInt32(&w.survivalNum))
}

// addSurvivalNum Added a new number of survivors
func (w *WorkerPool) addSurvivalNum() int {
	return int(atomic.AddInt32(&w.survivalNum, 1))
}

// addWorker Add a worker and work on tasks
func (w *WorkerPool) addWorker(task Task, core bool) error {
	var countSize int
	if core {
		countSize = w.coreSize
	} else {
		countSize = w.maxSize
	}
	if w.getSurvivalNum() > countSize {
		return errors.New("survivalNum is bigger than countSize")
	}
	if err := tool.LockWithSync(w.mu.Lock, w.mu.Unlock, func() error {
		if w.getSurvivalNum() > countSize {
			return errors.New("survivalNum is bigger than countSize")
		}
		// Create a worker
		worker, err := w.workerFactory.Create()
		if err != nil {
			return fmt.Errorf("create worker err:[%v]", err)
		}
		w.workerPool[worker] = struct{}{}
		w.addSurvivalNum()
		// Worker initialization
		return worker.Initialization()
	}); err != nil {
		return err
	}
	// The worker began to work
	return w.runWorker(task)
}

// runWorker worker processes the task
func (w *WorkerPool) runWorker(task Task) error {
	runnableTask := task
	var appropriateWorker Worker
	return w.safeGo.SafeGo(func() {
		// worker If you can't get the task, go to the destruction logic
		defer w.destroyWorker(appropriateWorker)
		// Get the most suitable worker (the least stressed)
		appropriateWorker = w.mostAppropriate()
		for {
			// If the delivered task is not empty, it is submitted
			if runnableTask != nil {
				appropriateWorker.Submit(runnableTask)
			}
			// If the task obtained by the queue is not empty, it is submitted
			if runnableTask = w.getTask(); runnableTask != nil {
				appropriateWorker.Submit(runnableTask)
			} else {
				return
			}
		}
	})
}

// mostAppropriate Get the most suitable worker
func (w *WorkerPool) mostAppropriate() Worker {
	selectionFun := GetSelection(w.workerSelection)
	if selectionFun != nil {
		return selectionFun(w.workerPool)
	}
	return nil
}

// idle Check whether the worker is idle
func (w *WorkerPool) idle() Worker {
	for worker := range w.workerPool {
		if worker.GetLoadFactor() == 0 {
			return worker
		}
	}
	return nil
}

// getTask Get a task from the queue
func (w *WorkerPool) getTask() (task Task) {
	wakeUp := time.After(w.nonCoreSurvivalTime)
	for {
		select {
		case getTask, ok := <-w.taskQueue:
			if !ok {
				return nil
			}
			return getTask
		case <-wakeUp:
			return nil
		}
	}
}

// destroyWorker Destroy a worker
func (w *WorkerPool) destroyWorker(worker Worker) {
	if worker == nil {
		return
	}
	// 任务队列中还有任务先不销毁
	if len(w.taskQueue) != 0 {
		return
	}
	worker.Destroy()
	w.delWorker(worker)
}

// delWorker Delete a worker
func (w *WorkerPool) delWorker(worker Worker) {
	_ = tool.LockWithSync(w.mu.Lock, w.mu.Unlock, func() error {
		delete(w.workerPool, worker)
		return nil
	})
}
